#自定义工具参数
from langchain_core.tools import StructuredTool
import asyncio

from pydantic import BaseModel,Field


class CalculatorInput(BaseModel):
    a:int = Field(description="第一个数字")
    b:int = Field(description="第二个数字")

def multiply(a: int, b: int) -> int:
    """同步乘积"""
    print("同步乘积")
    return a * b

async def asy_multiply(a: int, b: int) -> int:
    """异步乘积"""
    print("异步乘积")
    return a * b

async def main():
    # func传入同步方法，当你在同步上下文中调用工具时，它会使用这个同步方法来执行操作
    # coroutine传入异步方法，当你在异步上下文中调用工具时，它会使用这个异步方法来执行操作
    calculator = StructuredTool.from_function(
        func=multiply,
        coroutine=asy_multiply,
        name="calculator",
        description="乘积工具",
        args_schema=CalculatorInput,
        return_direct=True,
    )
    # 输出calculator的属性
    print(calculator.name)
    print(calculator.description)
    print(calculator.args)
    print(calculator.args_schema)
    # 同步调用
    print(calculator.invoke({"a": 2, "b": 3}))
    # 异步调用
    print(await calculator.ainvoke({"a": 9, "b": 10}))

# 执行异步main函数
asyncio.run(main())